import numpy as np

# 环境配置
ENV_SIZE = 12  # 12x12网格
INNER_SIZE = 10  # 10x10内部网格
NUM_COINS = 5  # 金币数量
NUM_OBSTACLES = 10  # 障碍物数量

# 单元格值
EMPTY = 0
START = 1
AGENT = 2
COIN = 3
OBSTACLE = 4
END = 5

# 动作
UP = 0
DOWN = 1
LEFT = 2
RIGHT = 3
STAY = 4

ACTION_NAMES = ["UP", "DOWN", "LEFT", "RIGHT", "STAY"]

# 方向向量 - 用于图神经网络
DIRECTIONS = [(-1, 0), (1, 0), (0, -1), (0, 1)]  # 上、下、左、右

# 每个单元格的可能状态数（用于one-hot编码）
CELL_FEATURES = 6

# 随机种子，保证可重现性
RANDOM_SEED = 42

# 模型配置
INPUT_SIZE = ENV_SIZE * ENV_SIZE * CELL_FEATURES  # 12*12*6 = 864
HIDDEN_SIZES = [512, 512, 256]
OUTPUT_SIZE = ENV_SIZE * ENV_SIZE * CELL_FEATURES  # 12*12*6 = 864

# GNN配置
GNN_NODE_FEATURES = 16   # 每个节点的特征数
GNN_EDGE_FEATURES = 4    # 每条边的特征数
GNN_HIDDEN_CHANNELS = 64 # GNN隐藏层通道数
GNN_NUM_LAYERS = 3       # GNN层数

# 智能体位置预测分支的增强容量
AGENT_BRANCH_HIDDEN_SIZES = [512, 256, 128, 64, 32]  # 增加层数和宽度

# 训练配置
BATCH_SIZE = 64
NUM_EPOCHS = 500  # 增加到500轮
LEARNING_RATE = 0.001
AGENT_LEARNING_RATE = 0.005  # 智能体分支使用更高的学习率
TRAIN_SPLIT = 0.7
VAL_SPLIT = 0.15
TEST_SPLIT = 0.15

# 正则化参数
DROPOUT_RATE = 0.3
WEIGHT_DECAY = 0.001

# 数据生成参数
NUM_SAMPLES = 10000
NUM_TRAJECTORIES = 50
MAX_STEPS_PER_TRAJECTORY = 200

# 损失权重 - 更关注智能体位置预测
ENV_LOSS_WEIGHT = 0.1  # 降低环境状态损失的权重
AGENT_LOSS_WEIGHT = 0.9  # 增加到0.9

# 多样性采样参数
MIN_DIVERSE_SAMPLES = 5000  # 最低要收集的多样性样本数
AGENT_MOVEMENT_RATIO = 0.8  # 确保80%的样本包含智能体移动

# 状态表示参数
USE_DISTANCE_MAPS = True    # 是否使用距离图作为额外特征
USE_DIRECTION_MAPS = True   # 是否使用方向图作为额外特征
POSITIONAL_ENCODING_DIM = 8 # 位置编码维度

# 早停参数
PATIENCE = 30  # 30轮无改善则停止
MIN_LR = 1e-6  # 最小学习率